package avlTree;

public class AVLTree {
    static class TreeNode{
        public int val;
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;
        public int bf; // 平衡因子

        public TreeNode(int val) {
            this.val = val;
        }
    }
    public TreeNode root;

    /**
     * 向AVL树中插入元素
     * @param val
     * @return
     */
    public boolean insert(int val){
        TreeNode node = new TreeNode(val);
        if(root == null){
            root = node;
            return true;
        }
        TreeNode temp = root;
        TreeNode parent = root;
        // 寻找插入的位置
        while (temp != null){
            parent = temp;
            if(node.val < temp.val){
                temp = temp.left;
            } else if(node.val == temp.val){
                return false;
            } else {
                temp = temp.right;
            }
        }
        // 此时temp == null
        if(node.val < parent.val){
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        temp = node;
        // 修改平衡因子
        while (parent != null){
            if(temp == parent.left){
                parent.bf--;
            } else {
                parent.bf++;
            }
            if(parent.bf == 0){
                break;// 已经平衡
            } else if (parent.bf == 1 || parent.bf == -1) {
                temp = parent;
                parent = temp.parent;
            } else {  // 左边高 右旋， 右边高 左旋
                if(parent.bf == 2){
                    if(temp.bf == 1){
                        rotateLeft(parent);
                    } else { // temp.bf == -1
                        rotateRightLeft(parent);
                    }
                } else { // parent.bf == -2
                    if(temp.bf == 1){ // 先左旋 再右旋
                        rotateLeftRight(parent);
                    } else {
                        rotateRight(parent);
                    }
                }
                break;
            } 
        }
        return true;
    }

    /**
     *  右单旋
     * @param parent
     */
    private void rotateRight(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        parent.left = subLR;
        subL.right = parent;
        if(subLR != null){
            subLR.parent = parent;
        }
        // 先记录parent的parent
        TreeNode pParent = parent.parent;
        parent.parent = subL;
        // 判断是否为根节点
        if(parent == root){
            root = subL;
            root.parent = null;
        } else { // 不是根节点时，判断是左子树还是右子树
            if(pParent.left == parent){
                pParent.left = subL;
            } else { // pParent.right = parent
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
        // 修改平衡因子
        subL.bf = 0;
        parent.bf = 0;
    }

    /**
     * 左单旋
     * @param parent
     */
    private void rotateLeft(TreeNode parent){
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        parent.right = subRL;
        subR.left = parent;
        if(subRL != null){
            subRL.parent = parent;
        }
        TreeNode pParent = parent.parent;
        parent.parent = subR;
        if (parent == root){
            root = subR;
            root.parent = null;
        } else {
            if(parent == pParent.left){
                pParent.left = subR;
            } else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        subR.bf = 0;
        parent.bf = 0;
    }

    /**
     * 左右双旋
     * @param parent
     */
    private void rotateLeftRight(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;
        rotateLeft(subL);
        rotateRight(parent);
        if(bf == -1){
            parent.bf = 1;
            subL.bf = 0;
        } else if(bf == 1){
            subL.bf = -1;
            parent.bf = 0;
        }
        subLR.bf = 0;
    }

    /**
     * 右左双旋
     * @param parent
     */
    private void rotateRightLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;
        rotateRight(subR);
        rotateLeft(parent);
        if(bf == -1){
            subR.bf = 1;
            parent.bf = 0;
        } else if(bf == 1){
            parent.bf = -1;
            subR.bf = 0;
        }
        subRL.bf = 0;
    }

    public void inorder(TreeNode root){
        if(root == null){
            return;
        }
        inorder(root.left);
        System.out.print(root.val + "  ");
        inorder(root.right);
    }

    public boolean isBalanced(TreeNode root){
        if(root == null){
            return true;
        }
        int left = height(root.left);
        int right = height(root.right);
        if(Math.abs(left - right) > 1){
            return false;
        }
        if(right - left != root.bf){
            System.out.println("节点" + root.val + " 错误");
            return false;
        }
        return isBalanced(root.left) && isBalanced(root.right);
    }

    private int height(TreeNode root){
        if(root == null){
            return 0;
        }
        int left = height(root.left);
        int right = height(root.right);
        return Math.max(left, right) + 1;
    }
    public static void main(String[] args) {
        AVLTree tree = new AVLTree();
        int[] arr1 = {16, 3, 7, 11, 9, 26, 18, 14, 15};
        int[] arr2 = {30, 20, 40, 35, 50, 60};
        for (int i = 0; i < arr2.length; i++) {
            tree.insert(arr2[i]);
        }
        tree.inorder(tree.root);
        System.out.println(tree.isBalanced(tree.root));
        System.out.println("+++++++++++++++");
    }
}
